import arrow
import time
import datetime
from celery import Celery
from celery.schedules import crontab
from kombu import Queue
from celery_once import QueueOnce
from flask import current_app
# from celery.schedules import crontab
from app.db.models.jobs import Job
from app.db.models.users import User
from app.db.models.balances import Balances
from app.db.models.asset_tasks import AssetTasks
from app.db.models.alertreceivers import AlertReceivers
from app.libs.oms import Client
from app.errors import OMSAPIException, OMSAPINotOk
from app.libs.sendcloud import SendCloud
from app.libs.yunpian import YunPian
from app.libs.enums import NotificationTargetQueue, TaskName
from app.config.settings import SITE_PORTRAIT_INTERVAL
from app.libs.utils import schedule_vul, handle_custom_urls_data


from celery.utils.log import get_task_logger

logger = get_task_logger(__name__)

celery_app = Celery(__name__)
celery_app.config_from_object(
    {
        "task_default_queue": NotificationTargetQueue.SCANV.value,
        "task_routes": {
            TaskName.LOAD_SCAN_EVENTS_FOR_EASM.value: {"queue": NotificationTargetQueue.EASM.value},
        }
    }
)


@celery_app.task(name="site_portrait_server")
def site_portrait_server():
    from app.celery.site_portraits import site_portraits
    site_portraits()


@celery_app.on_after_configure.connect
def setup_periodic_tasks(sender, **kwargs):
    # update_cyd_enabled
    sender.add_periodic_task(crontab(hour=19, minute=0), update_cyd_enabled.s(enable=False), name='update_cyd_enabled')
    sender.add_periodic_task(crontab(hour=f"*/{SITE_PORTRAIT_INTERVAL}", minute="5"), update_cyd_enabled.s(enable=True), name='update_cyd_enabled')

    # 网站画像服务
    sender.add_periodic_task(crontab(hour=f"*/{SITE_PORTRAIT_INTERVAL}", minute="55"), site_portrait_server.s(), name="site_portrait_server")

    # # Calls test('world') every 30 seconds
    # sender.add_periodic_task(30.0, test.s('world'), expires=10)

    # # Executes every Monday morning at 7:30 a.m.
    # sender.add_periodic_task(
    #     crontab(hour=7, minute=30, day_of_week=1),
    #     test.s('Happy Mondays!'),
    # )


@celery_app.task(base=QueueOnce, once={'graceful': False}, name='update_cyd_enabled')
def update_cyd_enabled(enable=False):
    logger.info('\n' * 5)
    logger.info('begin update_cyd_enabled, job_target__ipType=0')
    # 内部测试账号，有5w+资产，暂时不更新这部分资产状态
    except_uids = User.objects.filter(username__in=['zhichi@knownsec.com']).distinct('_id')
    uids = Balances.objects.filter(planExpireTime__gt=arrow.utcnow().datetime, uid__nin=except_uids).distinct('uid')
    query = {"target__ipType": 0, "uid__in": uids, "cydEnabled": enable}
    jobs = Job.objects.only(
        'targetUrl', 'uid', 'cydEnabled', 'target', 'vulSettings', 'enableMonitor', 'note'
    ).filter(**query)
    logger.info(f'domain jobs count: {jobs.count()}')
    oms_client = Client()
    sendcloud = SendCloud()
    yunpian = YunPian()

    true_count, false_count, fail_count = 0, 0, 0
    for job in jobs:
        user = User.objects.filter(pk=job.uid).first()
        time.sleep(0.2)  # 限制下频率
        try:
            resp = oms_client.get('/api/inner/asset_status/', params={
                'username': user.username,
                'asset': job.targetUrl
            }).get('data', False)
        except (OMSAPIException, OMSAPINotOk):
            fail_count += 1
            continue
        except Exception as e:
            fail_count += 1
            logger.exception(e)
            continue

        if resp is True:
            true_count += 1
        else:
            false_count += 1

        if job.cydEnabled != resp and isinstance(resp, bool):
            Job.objects(pk=job.pk).update_one(cydEnabled=resp)
            logger.info(
                f'target: {job.targetUrl}, status: {resp}, old:{job.cydEnabled}')

            # 如果resp==False 并且网站画像是开启的, 则关闭网站画像
            if not resp:
                unenable = False
                query = {
                    # "jobId": job.id,
                    "targetUrl": job.targetUrl,
                    "vulSettings__enable": True,
                    "vulSettings__enableSitePortraitTriggerMonitor": True
                }
                for asset_task in AssetTasks.objects.filter(**query):
                    asset_task.vulSettings.enableSitePortraitTriggerMonitor = False
                    asset_task.vulSettings.collect.includeUrl = []
                    asset_task.save()
                    schedule_vul(
                        job, asset_task.vulSettings,
                        prelogin=asset_task.prelogin.to_dict(),
                        custom_monitor_urls=handle_custom_urls_data(asset_task.customMonitorPage.to_dict()),
                        asset_task=asset_task,
                        alert_setting=asset_task.alertSettings,
                    )
                    unenable = True

                receiver = AlertReceivers.objects.filter(uid=job.uid, isDefault=True).first()
                # 发邮件短信提醒
                email = receiver.email
                phone = receiver.phone
                email_params = ('user', 'name', 'dz', 'url')
                sms_params = ('name', 'dz')
                detail = current_app.config.get('SCANVHOST')
                email_template = 'scanvmax_webportrait_has_stop'
                sms_template = '5238716'
                params = {'user': user.username, 'name': job.note, 'dz': job.targetUrl, 'url': detail}

                if receiver.emailVerified and email and unenable:
                    args = {}
                    for k in email_params:
                        args[k] = params.get(k, '')
                    args = {
                        'to': [email],
                        'sub': {f'%{k}%': [v] for k, v in args.items()}
                    }
                    status, result = sendcloud.sendmail(email_template, args)
                    logger.info(f"sendmail result {result}")

                if receiver.phoneVerified and phone and unenable:
                    args = {}
                    for k in sms_params:
                        args[k] = params.get(k, '')
                    args = '&'.join([f"#{k}#={v}" for k, v in args.items()])
                    status, res_msg = yunpian.tpl_send(phone, sms_template, args)
                    logger.info(f"sms result {res_msg}")

                #     # 总开关是开启的, 并且漏洞扫描是开启的, 则更新网站画像使用量
                #     if job.enableMonitor and job.vulSettings.enable:
                #         Balances.objects(uid=user.pk).update_one(
                #             dec__balance__sitePortrait__used=1)
                if unenable:
                    logger.info(
                        f'target: {job.targetUrl}, status: {resp}, disable site_port_rait')
    logger.info(f'domain jobs count: {jobs.count()}')
    logger.info(
        f'true_count: {true_count}, false_count: {false_count}, fail_count: {fail_count}')


@celery_app.task
def add(x, y):
    """
    加法
    :param x:
    :param y:
    :return:
    """
    return str(x + y)


@celery_app.task
def flask_app_context():
    """
    celery使用Flask上下文
    :return:
    """
    with current_app.app_context():
        return str(current_app.config)
